#  -*- coding: utf-8 -*-
from pymongo import UpdateOne, ASCENDING,DESCENDING
from util.stock_util import get_trading_dates, get_all_codes,multi_computer
from util.database import DB_CONN
import time
#import tushare as ts
from data.daily_crawler import DailyCrawler
import traceback
from datetime import datetime
"""
日K线数据的修复
"""


class DailyFixing:

    def multi_fixing(self,index=False,begin_date=None,end_date=None):
        if index:
            code_cursor = DB_CONN.stable.find(
                {"index": True},
                projection={'code': True, 'index': True, '_id': False})
            all_codes = [x['code'] for x in code_cursor]
        else:
            all_codes = get_all_codes()
        args = (index,begin_date,end_date)
        multi_computer(fill_data_at_suspension_days,all_codes,args)

    @staticmethod
    def fill_daily_k_at_suspension_days_at_date_one_collection(
            all_codes,index, all_dates, collection):
        """
        更新单个数据集的单个日期的数据
        :param basics:
        :param all_dates:
        :param collection:
        :return:
        """
        code_last_trading_daily_dict = dict()

        #按交易日期先后顺序遍历，记录当天能找到的数据，以备下一天查询，如果下一天停牌，用当天收盘价记录
        for date in all_dates:
            update_requests = []
            start_time =time.time()
            last_daily_code_set = set(code_last_trading_daily_dict.keys())
            for code in all_codes:

                trade_cursor = DB_CONN['stable'].find(
                    {'code': code,'index':index},
                    projection={'timeToMarket': True, '_id': False})

                # 如果循环日期小于上市日期
                try:
                    if date < trade_cursor[0]['timeToMarket']:
                        print('日期：%s, %s 还没上市，上市日期: %s' % (date, code, trade_cursor[0]["timeToMarket"]), flush=True)
                except:
                    print(f"error code:{code}")
                else:
                    # 找到当日数据
                    daily = DB_CONN[collection].find_one({'code': code, 'date': date, "index":index})
                    if daily is not None:
                        code_last_trading_daily_dict[code] = daily
                        last_daily_code_set.add(code)
                    else:
                        if code in last_daily_code_set:
                            last_trading_daily = code_last_trading_daily_dict[code]
                            suspension_daily_doc = {
                                'code': code,
                                'date': date,
                                'close': last_trading_daily['close'],
                                'open': last_trading_daily['close'],
                                'high': last_trading_daily['close'],
                                'low': last_trading_daily['close'],
                                'volume': 0,
                                'amount': 0.00,
                                'is_trading': False,
                                'index':index
                            }
                            update_requests.append(
                                UpdateOne(
                                    {'code': code, 'date': date},
                                    {'$set': suspension_daily_doc},
                                    upsert=True))
            if len(update_requests) > 0:
                update_result = DB_CONN[collection].bulk_write(update_requests, ordered=False)
                print('填充停牌数据，日期：%s，数据集：%s，插入：%4d条，更新：%4d条' %
                      (date, collection, update_result.upserted_count, update_result.modified_count), flush=True)
            end_time = time.time()
            print("遍历 %s 数据耗时：%.3f 秒" % (date,(end_time-start_time)))

    @staticmethod
    def fill_au_factor_pre_close(codes,begin_date, end_date):
        """
        为daily数据集填充：
        1. 复权因子au_factor，复权的因子计算方式：au_factor = hfq_close/close
        2. pre_close = close(-1) * au_factor(-1)/au_factor     (-1)表示前一天
        比如说：1.2是除权日，10送20
        date   close   pre_close  au_factor
        1.1    24       23           1
        1.2    8.4      8=24*（1/3)  3
        在1.2日，重新再计算1.1的前复权价格，就是24*（1/3)
        :param begin_date: 开始日期
        :param end_date: 结束日期
        """
        # 设置默认的日期范围,tdx有数据的时间

        for code in codes:
            start_time = time.time()
            hfq_daily_cursor = DB_CONN['daily_hfq'].find(
                {'code': code, 'date': {'$lte': end_date, '$gte': begin_date}, 'index': False},
                sort=[('date', ASCENDING)],
                projection={'date': True,  'close': True,'_id': False})

            date_hfq_close_dict = dict([(x['date'], x['close']) for x in hfq_daily_cursor])

            daily_cursor = DB_CONN['daily'].find(
                {'code': code, 'date': {'$lte': end_date, '$gte': begin_date}, 'index': False},
                sort=[('date', ASCENDING)],
                projection={'date': True, 'open': True, 'close': True,'_id': False}
            )

            last_close = -1
            last_au_factor = -1

            update_requests = []
            for daily in daily_cursor:
                date = daily['date']
                try:
                    close = daily['close']
                    open = daily['open']
                    doc = dict()

                    au_factor = round(date_hfq_close_dict[date]/close, 2)
                    doc['au_factor'] = au_factor
                    if last_close != -1 and last_au_factor != -1 and au_factor != 0 :
                        pre_close = last_close * last_au_factor / au_factor
                        # 2位四舍五入
                        doc['pre_close'] = round(pre_close, 2)
                    if au_factor == 0:
                        doc['pre_close'] = open
                    last_au_factor = au_factor
                    last_close = close

                    update_requests.append(
                        UpdateOne(
                            {'code': code, 'date': date, 'index': False},
                            {'$set': doc}))
                except:
                    print('计算复权因子时发生错误，股票代码：%s，日期：%s' % (code, date), flush=True)
                    # 恢复成初始值，防止用错
                    last_close = -1
                    last_au_factor = -1

            if len(update_requests) > 0:
                update_result = DB_CONN['daily'].bulk_write(update_requests, ordered=False)
                end_time = time.time()
                print('填充复权因子和前收，股票：%s，更新：%4d条, 耗时:  %.3f 秒' %
                      (code, update_result.modified_count,(end_time-start_time)), flush=True)

def fill_data_at_suspension_days(codes,index, begin_date=None, end_date=None):
    """
    :param codes:
    :param begin_date:
    :param end_date:
    :return:
    """
    df = DailyFixing()
    now = datetime.now()
    if begin_date is None:
        begin_date = "1990-12-19"

    # 结束日期默认为今天
    if end_date is None:
        end_date = now.strftime('%Y-%m-%d')

    all_dates = get_trading_dates(begin_date, end_date)
    if index:
        df.fill_daily_k_at_suspension_days_at_date_one_collection(
            codes,index, all_dates, 'daily')
    else:
        df.fill_daily_k_at_suspension_days_at_date_one_collection(
            codes,index, all_dates, 'daily')
        df.fill_daily_k_at_suspension_days_at_date_one_collection(
            codes,index, all_dates, 'daily_hfq')
        df.fill_daily_k_at_suspension_days_at_date_one_collection(
            codes,index, all_dates, 'daily_qfq')

        df.fill_au_factor_pre_close(codes,begin_date,end_date)


if __name__ == '__main__':
    #DailyFixing().fill_au_factor_pre_close("1995-02-24","2018-08-10")
    df = DailyFixing()
    df.multi_fixing(index=False, begin_date="2021-02-10",end_date="2021-02-18")
    #df.multi_fixing(index=True, begin_date=begin_date,end_date=current_date)

    #DailyFixing().fill_daily_k_at_suspension_days('1995-08-04', '1995-08-28')
    #DailyFixing().fill_is_trading_between('2018-07-26', '2018-08-07')
    # DailyFixing().fill_high_limit_low_limit('2015-01-01', '2018-12-31')
